# -*- coding: utf-8 -*-
"""
Created on Mon Oct 28 08:37:29 2024

@author: 51testing
"""
from prettytable import PrettyTable

def generate_table_mapping(model_class):
    mapping = {}
    for attribute_name in dir(model_class):
        # 过滤掉不需要的属性，比如内置属性和方法
        if not attribute_name.startswith('__') and not attribute_name.endswith('__'):
            attribute_value = getattr(model_class, attribute_name)
            # 检查属性是否是Mapped类型，并且是否是mapped_column的实例
            if str(type(attribute_value)) == "<class 'sqlalchemy.orm.attributes.InstrumentedAttribute'>":
                mapping[attribute_value.name] = attribute_name
    return mapping


def generate_pretty_table(res: list, pt_title = [], model_class = None):
    if not res:
        raise ValueError("No data to display")

    # 如果提供了model_class，则从模型中获取列名
    # if model_class is not None:
    #     if not issubclass(model_class, object):  # 检查model_class是否是一个类
    #         raise ValueError("model_class must be a class")
    #     pt_title = [x for x in model_class.__table__.columns.keys()]
    #     cm = generate_table_mapping(model_class)
    
    if hasattr(res[0], '_mapping') and hasattr(res[0], '__iter__'):
        dict_ = res[0]._mapping
        pt_title = []
        for i in dict_.keys():
            if isinstance(dict_[i], str) or isinstance(dict_[i], int):
                pt_title.append(i)
            else:
                for x in list(dict_[i].__mapper__.c.keys()):
                    count = 0
                    print(type(dict_[i]))
                    while (k := f"{dict_[i].__clsname__()}{str(count) if count else ''}.{x}") in pt_title:
                        count += 1
                    pt_title.append(k)
                    # count = 0
                    # k = f"{dict_[i].__clsname__()}.{x}"
                    # while k in pt_title:
                    #     count += 1
                    #     k = f"{dict_[i].__clsname__()}{count}.{x}"
                    # pt_title.append(k)
    else :
        pt_title = []
        for x in list(res[0].__mapper__.c.keys()):
            count = 0
            k = f"{res[0].__clsname__()}.{x}"
            while k in pt_title:
                count += 1
                k = f"{res[0].__clsname__()}{count}.{x}"
            pt_title.append(k)

        
    # 创建PrettyTable实例
    pt = PrettyTable(pt_title)

    # 迭代查询结果，添加到表格中
    for item in res:
        if hasattr(item, '_mapping'):
            dict_ = item._mapping
            row = []
            for i in dict_.keys():
                if isinstance(dict_[i], str) or isinstance(dict_[i], int):
                    row.append(dict_[i])
                else:
                    row.extend([getattr(dict_[i], x) for x in dict_[i].__mapper__.c.keys()])
        else:
            row = []
            row.extend([getattr(item, x) for x in item.__mapper__.c.keys()])
            
        # if isinstance(item, tuple):  # 如果item是元组
        #     row = [item[i] for i in range(len(pt_title))]
        # elif hasattr(item, '__iter__'):  # 如果item是可迭代的，但不是元组（例如列表或集合）
        #     row = list(item)
        # else :  # 如果item是单个对象
        #     row = [getattr(item, cm[x]) for x in pt_title]
        pt.add_row(row)
    
    print(pt)
    return pt